package com.atguigu.flink.datastream.transform;

import com.atguigu.flink.func.ClickSource;
import com.atguigu.flink.pojo.Event;
import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

/**
 * @author WEIYUNHUI
 * @date 2023/6/13 15:49
 *
 * 自定义分区:
 *    调用partitionCustom（Partitioner ， KeySelector）
 *
 *    Partitioner:
 *          int partition(K key, int numPartitions);  按照key来计算分区号
 *             K key : 当前数据的key
 *             int numPartitions : 下游的分区数
 *
 *   KeySelector:
 *          KEY getKey(IN value) throws Exception;   从数据中提取key
 *
 */
public class Flink06_UserDefinePartitioner {
    public static void main(String[] args) {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        //自定义分区，实现 中 外隔离
        DataStreamSource<Event> ds = env.addSource(new ClickSource());
        ds.partitionCustom(
                new Partitioner<String>() {
                    @Override
                    public int partition(String key, int numPartitions) {
                        if("Zhangs".equals(key) || "Lisi".equals(key)){
                            return 0 ;
                        }
                        return 1;
                    }
                }
                ,
                new KeySelector<Event, String>() {
                    @Override
                    public String getKey(Event value) throws Exception {
                        return value.getUser();
                    }
                }

        ).print().setParallelism(2);

        try {
            env.execute();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
}
